package snaptest.core;

import java.lang.annotation.Annotation;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

import org.junit.runners.BlockJUnit4ClassRunner;
import org.junit.runners.model.FrameworkMethod;
import org.junit.runners.model.InitializationError;
import org.junit.runners.model.Statement;

public class SnapTestRunner extends BlockJUnit4ClassRunner {

	public SnapTestRunner(Class<?> klass) throws InitializationError {
		super(klass);
	}

	@Override
	protected void validatePublicVoidNoArgMethods(
			Class<? extends Annotation> annotation, boolean isStatic,
			List<Throwable> errors) {
		List<Throwable> collectedErrors = new ArrayList<Throwable>();
		super.validatePublicVoidNoArgMethods(annotation, isStatic, collectedErrors);
		removeErrorsDueToReturnType(collectedErrors);
	}

	private List<Throwable> removeErrorsDueToReturnType(List<Throwable> collectedErrors) {
		Iterator<Throwable> iterator = collectedErrors.iterator();
		while(iterator.hasNext()) {
			Throwable error = iterator.next();
			if (error.getMessage().endsWith("should be void")) {
				iterator.remove();
			}
		}
		return collectedErrors;		
	}
	
	@Override
	protected Statement methodInvoker(FrameworkMethod method, Object test) {
		if (method.producesType(Object.class)) {
			MethodInvoker methodInvoker = createMethodInvoker(method, test);
			ReferenceFileLocator refLocator = createReferenceLocator(method, test);
			return new SnapTestMethodHandler(methodInvoker, refLocator);
		} else {
			return super.methodInvoker(method, test);
		}
	}

	protected MethodInvoker createMethodInvoker(FrameworkMethod method, Object test) {
		return new MethodInvoker(method, test);
	}

	protected ReferenceFileLocator createReferenceLocator(FrameworkMethod method, Object test) {
		return new ReferenceFileLocator(test.getClass(), method.getName());
	}

}
